import torch
import torch.nn as nn

# 测试评估函数
def eval(model, device, test_loader, criterion):
    # 切换到评估模式（关闭dropout/batchnorm等）
    model.eval()
    test_loss = 0
    correct = 0
    # 禁用梯度计算（节省内存）
    with torch.no_grad():
        for batch_idx, data, target in test_loader:
            """ 
                batch_idx: 数据批次
                data: 手写体图像
                target: 正确的分类值
            """
            data, target = data.to(device), target.to(device)
            print(f"\nEval Batch {batch_idx}")
            print(f"输入形状: {x.shape}，标签形状: {y.shape}")
            
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)  # 获取最大概率的类别
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    print(f'\nTest set: Average loss: {test_loss:.4f}, '
          f'Accuracy: {correct}/{len(test_loader.dataset)} ({100. * correct / len(test_loader.dataset):.2f}%)\n')
